package com.feidee.fd.sml.algorithm.component.ml.regression

import com.feidee.fd.sml.algorithm.component.ml.MLParam
import org.apache.spark.ml.PipelineStage
import org.apache.spark.ml.regression.AFTSurvivalRegression

/**
  * @Author: dongguosheng, songhaicheng
  * @Date: 2019/3/20 10:57
  * @Review songhaicheng
  * @Email: guosheng_dong@sui.com
  */
case class AFTSurvivalRegressionParam (
                                        override val input_pt: String,
                                        override val output_pt: String,
                                        override val hive_table: String,
                                        override val flow_time: String,
                                        override val featuresCol: String,
                                        override var labelCol: String,
                                        override var predictionCol: String,
                                        override val modelPath: String,
                                        override val metrics: Array[String],
                                        // 检查器列名，必填参数，默认 censor
                                        censorCol: String,
                                        // 分位数列 默认值 quantiles 如果设置该列，则会输出相应的分位数概率的分位数
                                        quantilesCol: String,
                                        // 分位数概率数组，值应在范围内（0，1），默认 [0.01, 0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.95, 0.99]
                                        quantileProbabilities: Array[Double],
                                        // 是否需要计算截距，默认 true
                                        fitIntercept: Boolean,
                                        // 最大迭代数，>= 0，默认 100
                                        maxIter: Int,
                                        // 优化算法迭代求解过程的收敛阀值，>= 0，默认 1E-6
                                        tol: Double,
                                        // Spark treeAggregate 算子的参数，适度增大可以加快训练速度，>= 2，默认 2
                                        aggregationDepth: Int
                                      ) extends MLParam {

  def this() = this(null, null, null, null, "features", "label", "prediction", null, new Array[String](0),
    "censor", "quantiles", Array(0.01, 0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.95, 0.99), true, 100, 1E-6, 2)

  override def verify(): Unit = {
    super.verify()
    require(quantileProbabilities.length > 0 &&
      quantileProbabilities.count(p => p > 0 && p < 1) == quantileProbabilities.length,
      "param quantileProbabilities can't be null and probability's range is (0, 1)")
    require(maxIter >= 0," param maxIter can't be negative")
    require(tol >= 0," param tol can't be negative")
    require(aggregationDepth >= 2," param aggregationDepth must be greater than 1")
  }

  override def toMap: Map[String, Any] = {
    var map = super.toMap
    map += ("censorCol" -> censorCol)
    map += ("quantilesCol" -> quantilesCol)
    map += ("quantileProbabilities" -> quantileProbabilities)
    map += ("fitIntercept" -> fitIntercept)
    map += ("maxIter" -> maxIter)
    map += ("tol" -> tol)
    map += ("aggregationDepth" -> aggregationDepth)
    map
  }
}


class AFTSurvivalRegressionComponent extends AbstractRegressionComponent[AFTSurvivalRegressionParam] {
  override def setUp(param: AFTSurvivalRegressionParam): PipelineStage = {
    val aft = new AFTSurvivalRegression()
      .setQuantilesCol(param.quantilesCol)
      .setAggregationDepth(param.aggregationDepth)
      .setCensorCol(param.censorCol)
      .setFeaturesCol(param.featuresCol)
      .setLabelCol(param.labelCol)
      .setPredictionCol(param.predictionCol)
      .setTol(param.tol)
      .setMaxIter(param.maxIter)
      .setFitIntercept(param.fitIntercept)
      .setQuantileProbabilities(param.quantileProbabilities)

    aft
  }

}

object AFTSurvivalRegressionComponent {
  def apply(paramStr: String): Unit = {
    new AFTSurvivalRegressionComponent()(paramStr)
  }

  def main(args: Array[String]): Unit = {
    AFTSurvivalRegressionComponent(args(0))
  }
}